import os
import sys
import json
import signal
import esptool 
import serial
import serial.tools.list_ports
import subprocess
from dataclasses import dataclass
from typing import Dict, Optional

DEFAULT_MONITOR_BAUD = 115200
DEFAULT_FLASH_BAUD = 468000

# ----------------- 现有 dataclass 配置 -----------------
@dataclass
class FlashSettings:
    flash_mode: str
    flash_size: str
    flash_freq: str

@dataclass
class FlashFileEntry:
    offset: str
    file: str
    encrypted: Optional[str] = None

@dataclass
class ExtraEsptoolArgs:
    after: str
    before: str
    stub: bool
    chip: str

@dataclass
class ESP32FlashConfig:
    flash_settings: FlashSettings
    flash_files: Dict[str, str]
    bootloader: FlashFileEntry
    partition_table: FlashFileEntry
    otadata: Optional[FlashFileEntry]
    app: FlashFileEntry
    extra_esptool_args: ExtraEsptoolArgs

def parse_config(data: dict) -> ESP32FlashConfig:
    return ESP32FlashConfig(
        flash_settings=FlashSettings(**data["flash_settings"]),
        flash_files=data["flash_files"],
        bootloader=FlashFileEntry(**data["bootloader"]),
        partition_table=FlashFileEntry(**data["partition-table"]),
        otadata=FlashFileEntry(**data["otadata"]) if "otadata" in data else None,
        app=FlashFileEntry(**data["app"]),
        extra_esptool_args=ExtraEsptoolArgs(**data["extra_esptool_args"])
    )

def load_flash_config(json_path: str = "") -> Optional[ESP32FlashConfig]:
    if not json_path:
        build_path = os.path.join(os.getcwd(), "build")
        json_path = os.path.join(build_path, "flasher_args.json")
        if not os.path.isdir(build_path) or not os.path.isfile(json_path):
            print("error: build/flasher_args.json not found")
            return None

    with open(json_path, "r", encoding="utf-8") as f:
        data = json.load(f)
    return parse_config(data)

# ----------------- ESP32 串口 & 烧录工具 -----------------
class ESP32Tool:
    def __init__(self, port, config: ESP32FlashConfig, default_monitor_baud=DEFAULT_MONITOR_BAUD):
        self.port = port
        self.ser = None
        self.config = config
        self.default_monitor_baud = default_monitor_baud

    def open_serial(self, baud):
        try:
            self.ser = serial.Serial(self.port, baudrate=baud, timeout=0.1)
            print(f"[INFO] Serial port {self.port} opened at {baud} baud.")
        except Exception as e:
            print(f"[ERROR] Failed to open serial port {self.port}: {e}")
            sys.exit(1)

    def close_serial(self):
        if self.ser and self.ser.is_open:
            self.ser.close()
            print(f"[INFO] Serial port {self.port} closed.")

    def monitor(self, baud=None):
        baud = baud or self.default_monitor_baud

        # --- 先复位芯片 ---
        try:
            cmd = [sys.executable, "-m", "esptool","--port", self.port,"--after", "hard-reset","--before", "default-reset","chip_id"]
            subprocess.run(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
            print("[INFO] Hard reset before entering monitor.")
        except Exception as e:
            print(f"[WARN] Reset failed before monitor: {e}")

        # --- 打开串口进入监控 ---
        try:
            self.open_serial(baud)
            print(f"[INFO] Entering monitor mode at {baud} baud. Press Ctrl+C to exit.")
            while True:
                if self.ser.in_waiting:
                    data = self.ser.read(self.ser.in_waiting)
                    print(data.decode(errors='ignore'), end='')
        except KeyboardInterrupt:
            print("\n[INFO] Exiting monitor function.")
        finally:
            try:
                self.close_serial()
            except KeyboardInterrupt:
                pass


    def flash(self, baud=None, json_path=""):
        try:
            baud = baud or DEFAULT_FLASH_BAUD

            # 确定 base_path, 用于解析相对路径
            if not json_path:
                build_path = os.path.join(os.getcwd(), "build")
                json_path = os.path.join(build_path, "flasher_args.json")
            base_path = os.path.dirname(json_path)

            # --------- 打印当前配置参数 ---------
            print("[INFO] ESP32FlashConfig 参数如下：")
            print(f"  Flash mode: {self.config.flash_settings.flash_mode}")
            print(f"  Flash size: {self.config.flash_settings.flash_size}")
            print(f"  Flash freq: {self.config.flash_settings.flash_freq}")
            print(f"  Bootloader: {self.config.bootloader}")
            print(f"  Partition table: {self.config.partition_table}")
            print(f"  OTADATA: {self.config.otadata}")
            print(f"  App: {self.config.app}")
            print(f"  Extra esptool args: {self.config.extra_esptool_args}")

            # --------- 构建 esptool 命令 ---------
            cmd = [sys.executable, "-m", "esptool","--port", self.port,"--baud", str(baud),"--after", "hard-reset","--before", "default-reset","write-flash"]


            # 按顺序写入 flash 文件, 只用 entries, 避免重复
            entries = [
                self.config.bootloader,
                self.config.partition_table,
                self.config.otadata,
                self.config.app
            ]
            for entry in entries:
                if entry and entry.file:
                    file_path = os.path.abspath(os.path.join(base_path, entry.file))
                    if os.path.isfile(file_path):
                        cmd.extend([entry.offset, file_path])
                    else:
                        print(f"[WARN] 文件不存在或路径错误: {file_path}")

            if len(cmd) <= 5:
                print("[ERROR] 没有有效的 flash 文件, 无法执行烧录。请检查 flasher_args.json")
                return

            print(f"[INFO] 运行 flash 命令: {' '.join(cmd)}")
            subprocess.run(cmd)
        
        except KeyboardInterrupt:
            print("\n[INFO] Exiting flash function.")

    def erase(self, baud=None):
        try:
            baud = baud or DEFAULT_FLASH_BAUD
            cmd = [sys.executable, "-m", "esptool", "--port", self.port, "--baud", str(baud), "erase-flash"]

            print(f"[INFO] 运行 erase 命令: {' '.join(cmd)}")
            subprocess.run(cmd)

        except KeyboardInterrupt:
            print("\n[INFO] Exiting erase function.")
        
    def info(self):
        try:
            print("[INFO] ESP32FlashConfig 参数如下：")
            print(f"  Flash mode: {self.config.flash_settings.flash_mode}")
            print(f"  Flash size: {self.config.flash_settings.flash_size}")
            print(f"  Flash freq: {self.config.flash_settings.flash_freq}")
            print(f"  Bootloader: {self.config.bootloader}")
            print(f"  Partition table: {self.config.partition_table}")
            print(f"  OTADATA: {self.config.otadata}")
            print(f"  App: {self.config.app}")
            print(f"  Extra esptool args: {self.config.extra_esptool_args}")

        except KeyboardInterrupt:
            print("\n[INFO] Exiting info function.")
            
# ----------------- 命令行交互 -----------------
def main():

    if len(sys.argv) < 2:
        print("Usage: python esp_tool.py COMx")
        sys.exit(1)
    port = sys.argv[1]

    default_mon_baud = DEFAULT_MONITOR_BAUD

    # 检查串口是否存在
    ports = [p.device for p in serial.tools.list_ports.comports()]
    if port not in ports:
        print(f"[ERROR] Port {port} not found.")
        sys.exit(1)

    # 检查全局 -b/--baud 参数（作用于 monitor 默认值）
    if "-b" in sys.argv:
        idx = sys.argv.index("-b")
        if idx+1 < len(sys.argv):
            default_mon_baud = int(sys.argv[idx+1])
    elif "--baud" in sys.argv:
        idx = sys.argv.index("--baud")
        if idx+1 < len(sys.argv):
            default_mon_baud = int(sys.argv[idx+1])

    config = load_flash_config()
    if not config:
        sys.exit(1)

    tool = ESP32Tool(port, config, default_monitor_baud=default_mon_baud)
    
    while True:
        try:
            cmd_line = input(">>> ").strip()
        except KeyboardInterrupt:
            print() 
            print("please enter \"exit\" to exit")
            continue
        except EOFError:
            print() 
            continue

        if not cmd_line:
            continue

        if cmd_line.lower() == "exit":
            print("[INFO] Exiting script.")
            break
        elif "flash" in cmd_line and "monitor" in cmd_line:
            # 处理 flash + monitor 连续命令
            flash_part = cmd_line.split("monitor")[0].strip()
            monitor_part = "monitor " + cmd_line.split("monitor")[1].strip()
            # 先 flash
            baud_flash = None
            parts = flash_part.split()
            if "--baud" in parts:
                baud_flash = int(parts[parts.index("--baud")+1])
            elif "-b" in parts:
                baud_flash = int(parts[parts.index("-b")+1])
            tool.flash(baud_flash)
            # 再 monitor
            baud_mon = None
            parts = monitor_part.split()
            if "--baud" in parts:
                baud_mon = int(parts[parts.index("--baud")+1])
            elif "-b" in parts:
                baud_mon = int(parts[parts.index("-b")+1])
            tool.monitor(baud_mon)
        elif cmd_line.startswith("monitor"):
            baud = None
            parts = cmd_line.split()
            if "--baud" in parts:
                baud = int(parts[parts.index("--baud")+1])
            elif "-b" in parts:
                baud = int(parts[parts.index("-b")+1])
            tool.monitor(baud)
        elif cmd_line.startswith("flash"):
            baud = None
            parts = cmd_line.split()
            if "--baud" in parts:
                baud = int(parts[parts.index("--baud")+1])
            elif "-b" in parts:
                baud = int(parts[parts.index("-b")+1])
            tool.flash(baud)
        elif cmd_line.startswith("erase"):
            baud = None
            parts = cmd_line.split()
            if "--baud" in parts:
                baud = int(parts[parts.index("--baud")+1])
            elif "-b" in parts:
                baud = int(parts[parts.index("-b")+1])
            tool.erase(baud)
        elif cmd_line.startswith("info"):
            tool.info()
        elif "help" in cmd_line or "?" in cmd_line:
            print("可用命令：")
            print("  monitor [--baud <baud>] / [-b <baud>]   打开串口监控 (默认 115200)")
            print("  flash   [--baud <baud>] / [-b <baud>]   烧录 ESP32 固件 (默认 468000)")
            print("  erase   [--baud <baud>] / [-b <baud>]   擦除整个 Flash (默认 468000)")
            print("  flash monitor                           先烧录, 再打开串口监控")
            print("  info                                    打印相关烧录参数")
            print("  exit                                    退出工具")
        else:
            print("[WARN] Unknown command. Plese enter \"help\" or \"?\" to get useful command.")

if __name__ == "__main__":
    main()
